import contextlib
import pathlib
import random
import os
import ruamel.yaml
import json
import io
from typing import Dict, Union, List, Optional, Generator
from copy import deepcopy

os.sep = '/'
yaml = ruamel.yaml.YAML()
yaml.default_style = '"'


def saveyaml(datadict: Dict[str, Dict[str, List[str]]], path: str) -> None:
    buf = io.StringIO()
    yaml.dump(datadict, buf)
    string = buf.getvalue()
    pathlib.Path(path).parent.mkdir(parents=True, exist_ok=True)
    with open(path, 'w') as f:
        f.seek(0)
        f.truncate()
        f.write(string)


def savejson(datadict: Dict, path: str) -> None:
    datadict = json.dumps(datadict, indent=4, sort_keys=False)
    with open(path, 'w') as f:
        f.seek(0)
        f.truncate()
        f.write(datadict)


def generatey(
    train_val_path: str,
    test_path: str,
    mask_patterns: Union[List, str] = None,
    img_pattern: Optional[str] = '*CT*.gz',
    yaml_savepath: Optional[str] = 'dataset.yaml',
) -> None:
    """generate dataset original yaml

    Args:
        train_val_path (str): Path of train & val set. e.g,
            '/home/yusongli/_dataset/_IIPL/ShuaiWang/20211214/shidaoai/sichuan'
        test_path (str): Path of test set. e.g,
            '/home/yusongli/_dataset/_IIPL/ShuaiWang/20211214/shidaoai/beijing'
    """
    if mask_patterns is None:
        mask_patterns = ['*T[_1 ]*.gz']
    if isinstance(mask_patterns, str):
        mask_patterns = [mask_patterns]
    path = pathlib.Path(train_val_path)
    data = []
    for pattern in mask_patterns:
        data.extend(
            {'image': list(item.parent.rglob(img_pattern))[0].as_posix(), 'label': item.as_posix()}
            for item in list(path.rglob(pattern))
        )
    random.seed(0)
    # Process sichuan data. Randomly split (train/val == 8/2)
    data_train = random.sample(data, int(len(data) * 0.8))
    data_val = [item for item in data if item not in data_train]
    yaml_dict = {'training': data_train, 'validation': data_val, 'test': []}
    with contextlib.suppress(Exception):
        path_test = pathlib.Path(test_path)
        for pattern in mask_patterns:
            yaml_dict['test'].extend(
                {'image': list(item.parent.rglob(img_pattern))[0].as_posix(), 'label': item.as_posix()}
                for item in list(path_test.rglob(pattern))
            )
    saveyaml(yaml_dict, yaml_savepath)


def filt_path(
    pathstr: Union[pathlib.Path, str], patterns: Optional[Union[List[str], str]] = None, echo: Optional[bool] = True
) -> Union[List[str], str]:
    """
    Find some target files in a specific folder path.
    Args:
        pathstr (str): A specific folder path.
        patterns (list(str)|str): The patterns you want to find.
            e.g, patterns=['*T[_1 ]*.gz']
    Return:
        target_files (list): A list of target files.
    """
    if patterns is None:
        return pathstr
    path = pathlib.Path(pathstr)
    if isinstance(patterns, str):
        patterns = [patterns]
    if patterns and isinstance(patterns, list):
        target_files = []
        for pattern in patterns:
            target_files.extend([item.as_posix() for item in list(path.rglob(pattern))])
        if not target_files:
            if echo:
                print(f"Warning: {pathstr} not catched.")
            return None
        return target_files if len(target_files) > 1 else target_files[0]


def exists(pathlist: Union[List[str], str]) -> None:
    '''
    Args:
        pathlist: path list
    '''
    if isinstance(pathlist, str):
        pathlist = [pathlist]
    for path in pathlist:
        assert pathlib.Path(path).exists(), path
    return pathlist[0] if len(pathlist) == 1 else pathlist


class YamlMetaGenerator(object):
    '''
    [
        [img, msk],
        ...
    ]
    '''

    def __init__(self, yaml_metadata: str, tags: List[str], keys: List[str]) -> None:
        with open(yaml_metadata, 'r') as y:
            dataset = yaml.load(y)

        self.tags = tags
        self.keys = keys
        self._yieldset = []
        for tag in self.tags:
            self._yieldset.extend(dataset[tag])

    def __len__(self) -> int:
        return len(self._yieldset)

    def __iter__(self) -> Generator[List[List[str]], None, None]:
        for item in self._yieldset:
            yield [item[key] for key in self.keys]


def j_to_y(jsonpath: str, yamlpath: str) -> None:
    """convert json datalist to yaml datalist

    Args:
        jsonpath (str): json path
        yamlpath (str): yaml path
    """
    with open(jsonpath, 'r') as f1:
        saveyaml(json.load(f1), yamlpath)


def replace(strings: Union[List[str], str], repdict: Dict[str, str]) -> str:
    '''
    Args:
        string (str|list):
        repdict: replace dict
    '''
    if isinstance(strings, str):
        strings = [strings]
    for i in range(len(strings)):
        if not isinstance(strings[i], str):
            continue
        for k, v in repdict.items():
            if isinstance(v, list):
                strings[i] = strings[i].replace(k, v[i])
            else:
                strings[i] = strings[i].replace(k, v)
    return strings[0] if len(strings) == 1 else strings


def set_pathy(yaml_path: str, yaml_savepath: Optional[str] = None, json_path: Optional[str] = None) -> None:
    """Set path from a base yaml file and save it into a new yaml for converting all kinds of path
    Args:
        yaml_path: dataset.yaml path. e.g, 'dataset/yaml/dataset.yaml'
        yaml_save_path: new dataset.yaml path. e.g, 'dataset/convert_dataset.yaml'
        common_root: common root path. e.g, /home/yusongli/_dataset/_IIPL/_Shuaiwang/20211223
    """
    with open(yaml_path, 'r') as y:
        dataset = yaml.load(y)
    # ! <<< open debug yusongli
    try:
        with open(json_path, 'r') as j:
            s_to_n = json.load(j)
    except Exception:
        s_to_n = None
    # ! >>> clos debug

    for tag in dataset.keys():
        # ! <<< open debug yusongli
        l = 's' if tag == 'test' else 'r'
        sps = pathlib.Path(f'/home/yusongli/_dataset/shidaoai/img/_out/nn/DATASET/nnUNet_raw_data_base/nnUNet_raw_data/Task602_Z2/start_positions_T{l}.txt')
        with open(sps, 'r') as f:
            sps = f.read()
        sps = sps.split('\n')
        # ! >>> clos debug
        for item in dataset[tag]:

            # s_to_n['training']['sichuan']['liutong']['306865']['nnformer'][1]

            img = pathlib.Path(item['image'])
            msk = pathlib.Path(item['label'])

            where = img.parents[2].name
            who = img.parents[1].name
            number = img.parents[0].name

            nn = s_to_n[tag][where][who][number]['nnformer'][1].split('_')[1].split('.')[0]

            item['image'] = []
            item['label'] = []

            # ! <<< open debug yusongli
            # roi_folder = os.path.join(
            #     '/home/yusongli/_dataset/shidaoai/img/_out',
            #     'wangqifeng-spacial-precropped_96_224_224-dilated_k7_i2-net_myunetr_val-roi_zoom',
            #     f'20220407/56/zoom-1.1/{where}/{who}/{number}'
            # )
            # roi_folder = pathlib.Path(roi_folder)
            # roi = [int(item) for item in next(roi_folder.iterdir()).name.split('_')]
            # ! ===
            sp = sps[int(nn)].split('|')[1].split('+')
            sp = map(lambda x: int(x), sp)
            # ! >>> clos debug

            pairs = (('image', img), ('label', msk))
            for pair in pairs:
                item[pair[0]].append(tag)
                item[pair[0]].append(pair[0])
                item[pair[0]].append(where)
                item[pair[0]].append(who)
                item[pair[0]].append(number)
                item[pair[0]].append(pair[1].name)
                item[pair[0]].append(nn)
                # ! <<< open debug yusongli
                # item[pair[0]].append(roi)
                # ! ===
                item[pair[0]].append(list(deepcopy(sp)))
                # ! >>> clos debug
            del pairs
    saveyaml(dataset, yaml_savepath)


if __name__ == '__main__':
    yaml_path = '/home/yusongli/Documents/shidaoai_new_project/data/train1332_val332_test264.yaml'
    yaml_convpath = '/home/yusongli/Documents/shidaoai_new_project/data/train1332_val332_test264_convert.yaml'
    yaml_metapath = '/home/yusongli/Documents/shidaoai_new_project/data/meta_data.yaml'
    yaml_metapath2 = '/home/yusongli/Documents/shidaoai_new_project/data/meta_data2.yaml'
    yaml_metapath3 = '/home/yusongli/Documents/shidaoai_new_project/data/meta_data3.yaml'
    s_to_n = '/home/yusongli/_dataset/shidaoai/img/_out/nn/s_to_n.json'

    set_pathy(yaml_path, yaml_metapath3, json_path=s_to_n)
    generator = YamlMetaGenerator(yaml_metapath2, tags=['training'], keys=['image', 'label'])
